home *** CD-ROM | disk | FTP | other *** search
/ EuroCD 3 / EuroCD 3.iso / Programming / Python1.4_Source / Lib / Complex.py < prev    next >
Text File  |  1998-06-24  |  7KB  |  276 lines
  1. # Complex numbers
  2. # ---------------
  3.  
  4. # This module represents complex numbers as instances of the class Complex.
  5. # A Complex instance z has two data attribues, z.re (the real part) and z.im
  6. # (the imaginary part).  In fact, z.re and z.im can have any value -- all
  7. # arithmetic operators work regardless of the type of z.re and z.im (as long
  8. # as they support numerical operations).
  9. #
  10. # The following functions exist (Complex is actually a class):
  11. # Complex([re [,im]) -> creates a complex number from a real and an imaginary part
  12. # IsComplex(z) -> true iff z is a complex number (== has .re and .im attributes)
  13. # Polar([r [,phi [,fullcircle]]]) ->
  14. #    the complex number z for which r == z.radius() and phi == z.angle(fullcircle)
  15. #    (r and phi default to 0)
  16. #
  17. # Complex numbers have the following methods:
  18. # z.abs() -> absolute value of z
  19. # z.radius() == z.abs()
  20. # z.angle([fullcircle]) -> angle from positive X axis; fullcircle gives units
  21. # z.phi([fullcircle]) == z.angle(fullcircle)
  22. #
  23. # These standard functions and unary operators accept complex arguments:
  24. # abs(z)
  25. # -z
  26. # +z
  27. # not z
  28. # repr(z) == `z`
  29. # str(z)
  30. # hash(z) -> a combination of hash(z.re) and hash(z.im) such that if z.im is zero
  31. #            the result equals hash(z.re)
  32. # Note that hex(z) and oct(z) are not defined.
  33. #
  34. # These conversions accept complex arguments only if their imaginary part is zero:
  35. # int(z)
  36. # long(z)
  37. # float(z)
  38. #
  39. # The following operators accept two complex numbers, or one complex number
  40. # and one real number (int, long or float):
  41. # z1 + z2
  42. # z1 - z2
  43. # z1 * z2
  44. # z1 / z2
  45. # pow(z1, z2)
  46. # cmp(z1, z2)
  47. # Note that z1 % z2 and divmod(z1, z2) are not defined,
  48. # nor are shift and mask operations.
  49. #
  50. # The standard module math does not support complex numbers.
  51. # (I suppose it would be easy to implement a cmath module.)
  52. #
  53. # Idea:
  54. # add a class Polar(r, phi) and mixed-mode arithmetic which
  55. # chooses the most appropriate type for the result:
  56. # Complex for +,-,cmp
  57. # Polar   for *,/,pow
  58.  
  59.  
  60. import types, math
  61.  
  62. if not hasattr(math, 'hypot'):
  63.     def hypot(x, y):
  64.         # XXX I know there's a way to compute this without possibly causing
  65.         # overflow, but I can't remember what it is right now...
  66.         return math.sqrt(x*x + y*y)
  67.     math.hypot = hypot
  68.  
  69. twopi = math.pi*2.0
  70. halfpi = math.pi/2.0
  71.  
  72. def IsComplex(obj):
  73.     return hasattr(obj, 're') and hasattr(obj, 'im')
  74.  
  75. def Polar(r = 0, phi = 0, fullcircle = twopi):
  76.     phi = phi * (twopi / fullcircle)
  77.     return Complex(math.cos(phi)*r, math.sin(phi)*r)
  78.  
  79. class Complex:
  80.  
  81.     def __init__(self, re=0, im=0):
  82.         if IsComplex(re):
  83.             im = im + re.im
  84.             re = re.re
  85.         if IsComplex(im):
  86.             re = re - im.im
  87.             im = im.re
  88.         self.re = re
  89.         self.im = im
  90.  
  91.     def __setattr__(self, name, value):
  92.         if hasattr(self, name):
  93.             raise TypeError, "Complex numbers have set-once attributes"
  94.         self.__dict__[name] = value
  95.  
  96.     def __repr__(self):
  97.         if not self.im:
  98.             return 'Complex(%s)' % `self.re`
  99.         else:
  100.             return 'Complex(%s, %s)' % (`self.re`, `self.im`)
  101.  
  102.     def __str__(self):
  103.         if not self.im:
  104.             return `self.re`
  105.         else:
  106.             return 'Complex(%s, %s)' % (`self.re`, `self.im`)
  107.  
  108.     def __coerce__(self, other):
  109.         if IsComplex(other):
  110.             return self, other
  111.         return self, Complex(other)    # May fail
  112.  
  113.     def __cmp__(self, other):
  114.         return cmp(self.re, other.re) or cmp(self.im, other.im)
  115.  
  116.     def __hash__(self):
  117.         if not self.im: return hash(self.re)
  118.         mod = sys.maxint + 1L
  119.         return int((hash(self.re) + 2L*hash(self.im) + mod) % (2L*mod) - mod)
  120.  
  121.     def __neg__(self):
  122.         return Complex(-self.re, -self.im)
  123.  
  124.     def __pos__(self):
  125.         return self
  126.  
  127.     def __abs__(self):
  128.         return math.hypot(self.re, self.im)
  129.         ##return math.sqrt(self.re*self.re + self.im*self.im)
  130.  
  131.  
  132.     def __int__(self):
  133.         if self.im:
  134.             raise ValueError, "can't convert Complex with nonzero im to int"
  135.         return int(self.re)
  136.  
  137.     def __long__(self):
  138.         if self.im:
  139.             raise ValueError, "can't convert Complex with nonzero im to long"
  140.         return long(self.re)
  141.  
  142.     def __float__(self):
  143.         if self.im:
  144.             raise ValueError, "can't convert Complex with nonzero im to float"
  145.         return float(self.re)
  146.  
  147.     def __nonzero__(self):
  148.         return not (self.re == self.im == 0)
  149.  
  150.     abs = radius = __abs__
  151.  
  152.     def angle(self, fullcircle = twopi):
  153.         return (fullcircle/twopi) * ((halfpi - math.atan2(self.re, self.im)) % twopi)
  154.  
  155.     phi = angle
  156.  
  157.     def __add__(self, other):
  158.         return Complex(self.re + other.re, self.im + other.im)
  159.  
  160.     __radd__ = __add__
  161.  
  162.     def __sub__(self, other):
  163.         return Complex(self.re - other.re, self.im - other.im)
  164.  
  165.     def __rsub__(self, other):
  166.         return Complex(other.re - self.re, other.im - self.im)
  167.  
  168.     def __mul__(self, other):
  169.         return Complex(self.re*other.re - self.im*other.im,
  170.                        self.re*other.im + self.im*other.re)
  171.  
  172.     __rmul__ = __mul__
  173.  
  174.     def __div__(self, other):
  175.         # Deviating from the general principle of not forcing re or im
  176.         # to be floats, we cast to float here, otherwise division
  177.         # of Complex numbers with integer re and im parts would use
  178.         # the (truncating) integer division
  179.         d = float(other.re*other.re + other.im*other.im)
  180.         if not d: raise ZeroDivisionError, 'Complex division'
  181.         return Complex((self.re*other.re + self.im*other.im) / d,
  182.                        (self.im*other.re - self.re*other.im) / d)
  183.  
  184.     def __rdiv__(self, other):
  185.         return other / self
  186.  
  187.     def __pow__(self, n, z=None):
  188.         if z is not None:
  189.             raise TypeError, 'Complex does not support ternary pow()'
  190.         if IsComplex(n):
  191.             if n.im: raise TypeError, 'Complex to the Complex power'
  192.             n = n.re
  193.         r = pow(self.abs(), n)
  194.         phi = n*self.angle()
  195.         return Complex(math.cos(phi)*r, math.sin(phi)*r)
  196.     
  197.     def __rpow__(self, base):
  198.         return pow(base, self)
  199.  
  200.  
  201. # Everything below this point is part of the test suite
  202.  
  203. def checkop(expr, a, b, value, fuzz = 1e-6):
  204.     import sys
  205.     print '       ', a, 'and', b,
  206.     try:
  207.         result = eval(expr)
  208.     except:
  209.         result = sys.exc_type
  210.     print '->', result
  211.     if (type(result) == type('') or type(value) == type('')):
  212.         ok = result == value
  213.     else:
  214.         ok = abs(result - value) <= fuzz
  215.     if not ok:
  216.         print '!!\t!!\t!! should be', value, 'diff', abs(result - value)
  217.  
  218.  
  219. def test():
  220.     testsuite = {
  221.         'a+b': [
  222.             (1, 10, 11),
  223.             (1, Complex(0,10), Complex(1,10)),
  224.             (Complex(0,10), 1, Complex(1,10)),
  225.             (Complex(0,10), Complex(1), Complex(1,10)),
  226.             (Complex(1), Complex(0,10), Complex(1,10)),
  227.         ],
  228.         'a-b': [
  229.             (1, 10, -9),
  230.             (1, Complex(0,10), Complex(1,-10)),
  231.             (Complex(0,10), 1, Complex(-1,10)),
  232.             (Complex(0,10), Complex(1), Complex(-1,10)),
  233.             (Complex(1), Complex(0,10), Complex(1,-10)),
  234.         ],
  235.         'a*b': [
  236.             (1, 10, 10),
  237.             (1, Complex(0,10), Complex(0, 10)),
  238.             (Complex(0,10), 1, Complex(0,10)),
  239.             (Complex(0,10), Complex(1), Complex(0,10)),
  240.             (Complex(1), Complex(0,10), Complex(0,10)),
  241.         ],
  242.         'a/b': [
  243.             (1., 10, 0.1),
  244.             (1, Complex(0,10), Complex(0, -0.1)),
  245.             (Complex(0, 10), 1, Complex(0, 10)),
  246.             (Complex(0, 10), Complex(1), Complex(0, 10)),
  247.             (Complex(1), Complex(0,10), Complex(0, -0.1)),
  248.         ],
  249.         'pow(a,b)': [
  250.             (1, 10, 1),
  251.             (1, Complex(0,10), 'TypeError'),
  252.             (Complex(0,10), 1, Complex(0,10)),
  253.             (Complex(0,10), Complex(1), Complex(0,10)),
  254.             (Complex(1), Complex(0,10), 'TypeError'),
  255.             (2, Complex(4,0), 16),
  256.         ],
  257.         'cmp(a,b)': [
  258.             (1, 10, -1),
  259.             (1, Complex(0,10), 1),
  260.             (Complex(0,10), 1, -1),
  261.             (Complex(0,10), Complex(1), -1),
  262.             (Complex(1), Complex(0,10), 1),
  263.         ],
  264.     }
  265.     exprs = testsuite.keys()
  266.     exprs.sort()
  267.     for expr in exprs:
  268.         print expr + ':'
  269.         t = (expr,)
  270.         for item in testsuite[expr]:
  271.             apply(checkop, t+item)
  272.     
  273.  
  274. if __name__ == '__main__':
  275.     test()
  276.